#!/usr/bin/env python3
"""
S3 Object Content-MD5 Tool - Content-MD5 校验演示工具

这个工具演示如何使用 S3 Content-MD5 校验功能进行数据完整性验证。

使用示例:
    python s3_object_content_md5.py --bucket my-bucket --key test-file.txt --file /path/to/file
    python s3_object_content_md5.py --bucket my-bucket --key test-file.txt --file /path/to/file --verify-md5
    python s3_object_content_md5.py --bucket my-bucket --key test-file.txt --file /path/to/file --invalid-md5
    python s3_object_content_md5.py --test-all --bucket my-bucket
"""

import argparse
import boto3
import json
import sys
import os
import hashlib
import base64
from datetime import datetime, timezone
from botocore.config import Config
from botocore.exceptions import ClientError, NoCredentialsError


class ContentMD5Client:
    """S3 Content-MD5 校验客户端"""
    
    def __init__(self, endpoint_url='http://localhost:8000', 
                 access_key='AKIAIOSFODNN7EXAMPLE', 
                 secret_key='wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY',
                 region='us-east-1'):
        """
        初始化 Content-MD5 校验客户端
        
        Args:
            endpoint_url: S3 服务端点
            access_key: AWS Access Key
            secret_key: AWS Secret Key
            region: AWS 区域
        """
        self.endpoint_url = endpoint_url
        self.access_key = access_key
        self.secret_key = secret_key
        self.region = region
        self._client = None
    
    def _get_client(self):
        """获取 S3 客户端"""
        if self._client is None:
            self._client = boto3.client(
                's3',
                endpoint_url=self.endpoint_url,
                aws_access_key_id=self.access_key,
                aws_secret_access_key=self.secret_key,
                region_name=self.region,
        config=Config(s3={'addressing_style': 'path'})
    )
        return self._client
    
    def _calculate_md5(self, data):
        """计算数据的MD5哈希值"""
        if isinstance(data, str):
            data = data.encode('utf-8')
        return hashlib.md5(data).hexdigest()
    
    def _calculate_md5_base64(self, data):
        """计算数据的MD5哈希值（Base64编码）"""
        if isinstance(data, str):
            data = data.encode('utf-8')
        md5_hash = hashlib.md5(data).digest()
        return base64.b64encode(md5_hash).decode('utf-8')
    
    def upload_with_md5(self, bucket, key, file_path, verify_md5=True, invalid_md5=False):
        """上传文件并验证Content-MD5"""
        try:
            # 检查文件是否存在
            if not os.path.exists(file_path):
                print(f"❌ 错误: 文件不存在: {file_path}")
                return False
            
            # 读取文件内容
            with open(file_path, 'rb') as f:
                file_data = f.read()
            
            file_size = len(file_data)
            print(f"📁 文件信息:")
            print(f"   路径: {file_path}")
            print(f"   大小: {file_size} 字节 ({file_size / 1024:.2f} KB)")
            
            # 计算MD5
            if verify_md5:
                if invalid_md5:
                    # 使用错误的MD5进行测试
                    content_md5 = 'deadbeefdeadbeefdeadbeefdeadbeef'
                    print(f"🧪 测试模式: 使用错误的MD5 - {content_md5}")
                else:
                    content_md5 = self._calculate_md5(file_data)
                    print(f"🔐 计算MD5: {content_md5}")
            else:
                content_md5 = None
                print("📤 上传模式: 不使用Content-MD5校验")
            
            # 上传文件
            upload_args = {
                'Bucket': bucket,
                'Key': key,
                'Body': file_data
            }
            
            if content_md5:
                upload_args['ContentMD5'] = content_md5
            
            print(f"\n📤 上传文件到 s3://{bucket}/{key}...")
            response = self._get_client().put_object(**upload_args)
            
            etag = response.get('ETag', '').strip('"')
            print(f"✅ 上传成功！")
            print(f"   🏷️  ETag: {etag}")
            
            if verify_md5 and not invalid_md5:
                # 验证ETag是否与MD5匹配
                if etag == content_md5:
                    print(f"   ✅ MD5校验通过")
                else:
                    print(f"   ⚠️  ETag与MD5不匹配")
            
            return True
            
        except ClientError as e:
            error_code = e.response.get('Error', {}).get('Code')
            error_message = e.response.get('Error', {}).get('Message')
            
            if error_code == 'BadDigest' and invalid_md5:
                print(f"✅ 预期结果: Content-MD5校验失败 [{error_code}]")
                print(f"   错误信息: {error_message}")
                return True
            else:
                print(f"❌ 上传失败 [{error_code}]: {error_message}")
                return False
        except Exception as e:
            print(f"❌ 上传失败: {e}")
            return False
    
    def test_content_md5_valid(self, bucket):
        """测试正确的Content-MD5校验"""
        print("🧪 测试1: 正确的Content-MD5校验")
        print("-" * 40)
        
        key = 'test-md5-valid.txt'
        test_data = b'Hello, Content-MD5 validation!'
        
        try:
            content_md5 = self._calculate_md5(test_data)
            print(f"📝 测试数据: {test_data.decode('utf-8')}")
            print(f"🔐 计算MD5: {content_md5}")
            
            response = self._get_client().put_object(
                Bucket=bucket,
            Key=key,
                Body=test_data,
                ContentMD5=content_md5
            )
            
            etag = response.get('ETag', '').strip('"')
            print(f"✅ 上传成功！ETag: {etag}")
            
            if etag == content_md5:
                print("✅ MD5校验通过")
        return True
            else:
                print("⚠️  ETag与MD5不匹配")
        return False

        except ClientError as e:
            error_code = e.response.get('Error', {}).get('Code')
            error_message = e.response.get('Error', {}).get('Message')
            print(f"❌ 测试失败 [{error_code}]: {error_message}")
            return False
        except Exception as e:
            print(f"❌ 测试失败: {e}")
            return False
    
    def test_content_md5_invalid(self, bucket):
        """测试错误的Content-MD5校验"""
        print("\n🧪 测试2: 错误的Content-MD5校验")
        print("-" * 40)
        
        key = 'test-md5-invalid.txt'
        test_data = b'Hello, Content-MD5 validation!'
    wrong_md5 = 'deadbeefdeadbeefdeadbeefdeadbeef'
    
    try:
            print(f"📝 测试数据: {test_data.decode('utf-8')}")
            print(f"🔐 使用错误MD5: {wrong_md5}")
            
            self._get_client().put_object(
                Bucket=bucket,
            Key=key,
                Body=test_data,
            ContentMD5=wrong_md5
        )
            
            print("❌ 测试失败: 应该被拒绝但成功了")
        return False
            
    except ClientError as e:
        error_code = e.response.get('Error', {}).get('Code')
            error_message = e.response.get('Error', {}).get('Message')
            
        if error_code == 'BadDigest':
                print(f"✅ 预期结果: Content-MD5校验失败 [{error_code}]")
                print(f"   错误信息: {error_message}")
            return True
        else:
                print(f"❌ 测试失败: 期望BadDigest错误，得到 [{error_code}]: {error_message}")
                return False
        except Exception as e:
            print(f"❌ 测试失败: {e}")
            return False

    def test_multipart_content_md5(self, bucket):
        """测试多部分上传的Content-MD5校验"""
        print("\n🧪 测试3: 多部分上传Content-MD5校验")
        print("-" * 40)
        
    key = 'test-multipart-md5.bin'
    
        try:
            # 创建多部分上传
            response = self._get_client().create_multipart_upload(Bucket=bucket, Key=key)
    upload_id = response['UploadId']
            print(f"🚀 创建多部分上传: {upload_id}")
    
            # 分片1：正确的MD5
        part1_data = b'A' * 5242880  # 5MB
            part1_md5 = self._calculate_md5(part1_data)
            print(f"📤 上传分片1 (5MB), MD5: {part1_md5}")
        
            part1_response = self._get_client().upload_part(
                Bucket=bucket,
            Key=key,
            PartNumber=1,
            UploadId=upload_id,
            Body=part1_data,
            ContentMD5=part1_md5
        )
            
            part1_etag = part1_response.get('ETag', '').strip('"')
            print(f"✅ 分片1上传成功！ETag: {part1_etag}")
        
            # 分片2：错误的MD5
        part2_data = b'B' * 5242880  # 5MB
        wrong_md5 = 'deadbeefdeadbeefdeadbeefdeadbeef'
            print(f"📤 上传分片2 (5MB), 使用错误MD5: {wrong_md5}")
        
        try:
                self._get_client().upload_part(
                    Bucket=bucket,
                Key=key,
                PartNumber=2,
                UploadId=upload_id,
                Body=part2_data,
                ContentMD5=wrong_md5
            )
                
                print("❌ 测试失败: 分片2应该被拒绝但成功了")
                self._get_client().abort_multipart_upload(Bucket=bucket, Key=key, UploadId=upload_id)
            return False
                
        except ClientError as e:
            error_code = e.response.get('Error', {}).get('Code')
                error_message = e.response.get('Error', {}).get('Message')
                
            if error_code == 'BadDigest':
                    print(f"✅ 预期结果: 分片2 Content-MD5校验失败 [{error_code}]")
                    print(f"   错误信息: {error_message}")
                    
                    # 清理：中止多部分上传
                    self._get_client().abort_multipart_upload(Bucket=bucket, Key=key, UploadId=upload_id)
                    print("🧹 已清理多部分上传")
                return True
            else:
                    print(f"❌ 测试失败: 期望BadDigest错误，得到 [{error_code}]: {error_message}")
                    self._get_client().abort_multipart_upload(Bucket=bucket, Key=key, UploadId=upload_id)
                return False
                
    except Exception as e:
            print(f"❌ 测试失败: {e}")
        return False

    def test_without_content_md5(self, bucket):
        """测试不使用Content-MD5的上传"""
        print("\n🧪 测试4: 不使用Content-MD5的上传")
        print("-" * 40)
        
    key = 'test-no-md5.txt'
        test_data = b'Hello, no MD5 validation!'
        
        try:
            print(f"📝 测试数据: {test_data.decode('utf-8')}")
            print("📤 上传模式: 不使用Content-MD5校验")
            
            response = self._get_client().put_object(
                Bucket=bucket,
            Key=key,
                Body=test_data
        )
            
            etag = response.get('ETag', '').strip('"')
            print(f"✅ 上传成功！ETag: {etag}")
        return True
            
    except ClientError as e:
            error_code = e.response.get('Error', {}).get('Code')
            error_message = e.response.get('Error', {}).get('Message')
            print(f"❌ 测试失败 [{error_code}]: {error_message}")
            return False
        except Exception as e:
            print(f"❌ 测试失败: {e}")
        return False

    def run_all_tests(self, bucket):
        """运行所有Content-MD5测试"""
        print("🚀 Content-MD5 校验测试套件")
        print("=" * 50)
        print(f"📦 测试桶: {bucket}")
        print(f"🌐 端点: {self.endpoint_url}")
        print("=" * 50)
    
    tests = [
            ("正确的Content-MD5校验", self.test_content_md5_valid),
            ("错误的Content-MD5校验", self.test_content_md5_invalid),
            ("多部分上传Content-MD5校验", self.test_multipart_content_md5),
            ("不使用Content-MD5的上传", self.test_without_content_md5)
    ]
    
    passed = 0
    total = len(tests)
    
        for test_name, test_func in tests:
            print(f"\n🧪 {test_name}")
            print("-" * 50)
            
            try:
                if test_func(bucket):
            passed += 1
                    print(f"✅ {test_name} - 通过")
                else:
                    print(f"❌ {test_name} - 失败")
            except Exception as e:
                print(f"❌ {test_name} - 异常: {e}")
        
        print("\n" + "=" * 50)
        print(f"📊 测试结果: {passed}/{total} 通过")
        
        if passed == total:
            print("🎉 所有Content-MD5测试通过！")
            return True
        else:
            print(f"⚠️  {total - passed} 个测试失败")
            return False


def parse_args():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(
        description='S3 Content-MD5 校验演示工具',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
使用示例:
  # 上传文件并验证MD5
  python s3_object_content_md5.py --bucket my-bucket --key test-file.txt --file /path/to/file --verify-md5
  
  # 上传文件不验证MD5
  python s3_object_content_md5.py --bucket my-bucket --key test-file.txt --file /path/to/file
  
  # 测试错误的MD5（应该失败）
  python s3_object_content_md5.py --bucket my-bucket --key test-file.txt --file /path/to/file --invalid-md5
  
  # 运行所有测试
  python s3_object_content_md5.py --test-all --bucket my-bucket
        """
    )
    
    # 基本参数
    parser.add_argument('--bucket', required=True, help='桶名称')
    parser.add_argument('--key', help='对象键名')
    parser.add_argument('--file', help='要上传的文件路径')
    parser.add_argument('--verify-md5', action='store_true', 
                       help='启用Content-MD5校验')
    parser.add_argument('--invalid-md5', action='store_true',
                       help='使用错误的MD5进行测试（应该失败）')
    parser.add_argument('--test-all', action='store_true',
                       help='运行所有Content-MD5测试')
    
    # 连接参数
    parser.add_argument('--endpoint', default='http://localhost:8000',
                       help='S3 服务端点 (默认: http://localhost:8000)')
    parser.add_argument('--access-key', default='AKIAIOSFODNN7EXAMPLE',
                       help='AWS Access Key (默认: 教学用凭据)')
    parser.add_argument('--secret-key', default='wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY',
                       help='AWS Secret Key (默认: 教学用凭据)')
    parser.add_argument('--region', default='us-east-1',
                       help='AWS 区域 (默认: us-east-1)')
    
    # 输出格式
    parser.add_argument('--format', choices=['simple', 'json'], default='simple',
                       help='输出格式 (默认: simple)')
    parser.add_argument('--verbose', action='store_true',
                       help='详细输出')
    
    return parser.parse_args()


def main():
    """主函数"""
    args = parse_args()
    
    # 创建客户端
    client = ContentMD5Client(
        endpoint_url=args.endpoint,
        access_key=args.access_key,
        secret_key=args.secret_key,
        region=args.region
    )
    
    try:
        if args.test_all:
            # 运行所有测试
            success = client.run_all_tests(args.bucket)
            sys.exit(0 if success else 1)
        
        else:
            # 单文件上传模式
            if not args.key:
                print("❌ 错误: 单文件上传模式需要 --key 参数")
                sys.exit(1)
            
            if not args.file:
                print("❌ 错误: 单文件上传模式需要 --file 参数")
                sys.exit(1)
            
            print("🚀 S3 Content-MD5 校验工具")
            print("=" * 50)
            print(f"📦 桶: {args.bucket}")
            print(f"🔑 键: {args.key}")
            print(f"🌐 端点: {args.endpoint}")
            print("=" * 50)
            
            success = client.upload_with_md5(
                bucket=args.bucket,
                key=args.key,
                file_path=args.file,
                verify_md5=args.verify_md5,
                invalid_md5=args.invalid_md5
            )
            
            if success:
                print("\n✅ 操作完成！")
    else:
                print("\n❌ 操作失败！")
                sys.exit(1)
        
    except KeyboardInterrupt:
        print("\n⚠️  操作被用户中断")
        sys.exit(1)
    except Exception as e:
        print(f"\n❌ 操作失败: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()